import torch

from problems.TREE import Tree


class TREE6(Tree):
    def __init__(self, var_dim: int, obj_dim: int, max_fun_eval, kwargs: dict):
        super().__init__(var_dim, obj_dim, max_fun_eval, 5, 12, self.init_dec_boundary2, 3, kwargs)

    def eval_value(self, pop_dec):
        # tensor2excel(pop_dec)
        # pop_dec = excel2tensor("2025-04-07_21-18-47.xlsx")
        # 获取输入张量的形状
        N = pop_dec.size(0)  # PopDec 的形状为 (N, ...)
        KP = self.data.size(1)  # obj.Data 的形状为 (T, KP)
        T, K = self.mean.size()  # obj.Mean 的形状为 (T, K)

        # 将 PopDec 重塑为 (N, T, K) 并加上 obj.Mean 的广播
        pop_dec = pop_dec.T.reshape(K, T, N).permute(2, 1, 0) + self.mean.T.reshape(K, T, 1).permute(2, 1, 0)

        # 初始化目标值 pop_obj
        pop_obj = torch.zeros(N, self.obj_dim, dtype=torch.double)

        # 第一个目标函数
        eA, eAA = super().calc_ea2(pop_dec, KP, T)
        # 对 eA 求和得到第一个目标值
        pop_obj[:, 0] = eA.sum(dim=(1, 2))  # 沿第 2 和第 3 维度求和

        # 第二个目标函数
        # 计算 Delta = std(eA[:, 2:] - eA[:, :-1])
        eA_diff = eA[:, 1:] - eA[:, :-1]  # 计算差分
        Delta = torch.std(eA_diff, dim=1, keepdim=True)  # 沿第 2 维度计算标准差

        # 对 Delta 求和得到第二个目标值
        pop_obj[:, 1] = Delta.permute(2, 1, 0).reshape(-1, N).T.sum(dim=1)  # 展平后求和

        # 计算 Delta = std(eA[:, 2:] - eA[:, :-1])
        eA_diff = eAA[:, 1:] - eAA[:, :-1]  # 计算差分
        Delta = torch.std(eA_diff, dim=1, keepdim=True)  # 沿第 2 维度计算标准差

        # 对 Delta 求和得到第二个目标值
        pop_obj[:, 2] = Delta.permute(2, 1, 0).reshape(-1, N).T.sum(dim=1)  # 展平后求和

        return pop_obj
